import { defineStore } from 'pinia';
import { store } from '/@/store';
import { WebSocketResult, useWebSocket } from '@vueuse/core';
import { sleep } from '/@/utils';
import { isArray } from '/@/utils/is';

type wsKey = 'telemetry' | 'notifications';

type WebSocketResultObj = {
  [key in wsKey]: WebSocketResult<any> | undefined;
};

interface WebsocketState {
  cmdId: number;
  websocket: WebSocketResultObj;
  callbackMap: Map<number, (data: any) => void>;
}

export const useWebsocketStore = defineStore({
  id: 'websocket-store',
  state: (): WebsocketState => ({
    cmdId: 0,
    websocket: {
      telemetry: undefined,
      notifications: undefined,
    },
    callbackMap: new Map(),
  }),
  getters: {
    getCmdId(): number {
      return this.cmdId;
    },
  },
  actions: {
    close(flag: wsKey[]): any {
      flag.forEach((item) => {
        if (this.websocket[item]) {
          this.websocket[item]?.close();
        }
        this.websocket[item] = undefined;
      });
    },
    getAndIncrementCmdId(): number {
      this.cmdId = this.cmdId + 1;
      return this.cmdId;
    },
    async send(
      cmdId: number | Array<number>,
      data: any,
      callback?: (data: any) => void,
      flag: wsKey = 'telemetry',
    ): Promise<boolean | undefined> {
      if (!this.websocket[flag]) {
        this.initWebsocket(flag);
      }
      if (this.websocket[flag]) {
        if (this.websocket[flag]?.status == 'CLOSED') {
          this.websocket[flag]?.open();
          await sleep(500);
        }
        if (callback) {
          if (isArray(cmdId)) {
            cmdId.forEach((i) => this.callbackMap.set(i, callback));
          } else {
            this.callbackMap.set(cmdId, callback);
          }
        }
        return this.websocket[flag]?.send(JSON.stringify(data));
      }
      return false;
    },
    async unsubscribe(cmdId: number | Array<number>, data: any, flag: wsKey = 'telemetry') {
      if (this.websocket[flag]) {
        if (this.websocket[flag]?.status == 'CLOSED') {
          this.websocket[flag]?.open();
          await sleep(500);
        }
        this.websocket[flag]?.send(JSON.stringify(data));
      }
      if (isArray(cmdId)) {
        cmdId.forEach((i) => this.callbackMap.delete(i));
      } else {
        this.callbackMap.delete(cmdId);
      }
    },
    initWebsocket(flag: wsKey): WebSocket | undefined {
      const useWebSocketReturn: any = useWebSocket(
        `ws://${import.meta.env.VITE_IP_PORT}/api/ws/plugins/${flag}?token=${localStorage.getItem('jwt_token')}`,
        {
          autoReconnect: false,
          autoClose: false,
          onMessage: this.onMessage,
        },
      );
      this.websocket[flag] = useWebSocketReturn;
      return this.websocket[flag]?.ws;
    },
    onMessage(ws: WebSocket, { data }: MessageEvent): any {
      const dataObj = JSON.parse(data);
      if (dataObj.hasOwnProperty('subscriptionId')) {
        this.callbackMap.get(dataObj.subscriptionId)?.(dataObj);
      } else if (dataObj.hasOwnProperty('cmdId')) {
        this.callbackMap.get(dataObj.cmdId)?.(dataObj);
      }
    },
  },
});

export function useWebsocketStoreWithOut() {
  return useWebsocketStore(store);
}
